#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
欧意（OKX）K线图实时显示 - 使用mplfinance库
支持获取合约历史K线数据并实时更新
"""

import os
import time
import json
import asyncio
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
from matplotlib.animation import FuncAnimation
import mplfinance as mpf
from typing import Dict, List, Tuple, Optional
import matplotlib.font_manager as fm
from matplotlib.widgets import MultiCursor
import threading
import matplotlib.dates as mdates
import pytz
import mplcursors  # 添加mplcursors库用于鼠标悬浮显示

from okx_market import OKXMarket

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'SimSun', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

class OKXMplfinanceChart:
    """
    使用mplfinance库显示OKX交易所的K线图
    """
    def __init__(self, symbol="BTC-USDT-SWAP", interval="1m", limit=100, proxy=None):
        """
        初始化K线图显示
        
        参数:
            symbol (str): 交易对，例如 "BTC-USDT-SWAP" 表示BTC永续合约
            interval (str): K线间隔，例如 "1m", "5m", "15m", "1H", "4H", "1D"
            limit (int): 获取历史K线的数量限制
            proxy (str): 代理服务器地址，例如 "http://localhost:10809"
        """
        self.symbol = symbol
        self.interval = interval
        self.limit = limit
        self.proxy = proxy
        
        # 创建OKX市场数据对象
        self.market = OKXMarket(proxy=proxy)
        
        # 存储K线数据的DataFrame
        self.df = None
        
        # 图表更新间隔（毫秒）
        self.update_interval = 2000
        
        # 动画对象
        self.ani = None
        
        # 图表和轴对象
        self.fig = None
        self.ax1 = None
        
        # 设置时区为UTC+8
        self.timezone = pytz.timezone('Asia/Shanghai')
        
        # 设置mplfinance样式
        self.style = mpf.make_mpf_style(
            marketcolors=mpf.make_marketcolors(
                up='red', down='green',
                edge='inherit',
                wick='inherit',

                
                volume='inherit'
            ),
            gridstyle='--',
            gridcolor='gray',
            rc={'font.family': 'SimHei'},  # 使用中文字体
            facecolor='white',  # 设置背景为白色
            figcolor='white',   # 设置图形背景为白色
            edgecolor='black',  # 设置边框颜色为黑色
            y_on_right=False    # 价格显示在左侧
        )
        
        # 技术指标颜色
        self.colors = ['blue', 'orange', 'purple', 'brown', 'pink']
        
        # 添加鼠标悬浮提示对象
        self.cursor = None
        
        # 添加K线对象列表，用于鼠标悬浮显示
        self.candle_plots = None
    
    def process_kline_data(self, data):
        """
        处理K线数据，转换为DataFrame格式
        
        参数:
            data (list): K线数据列表
            
        返回:
            pandas.DataFrame: 处理后的K线数据
        """
        # 创建DataFrame
        columns = ['timestamp', 'open', 'high', 'low', 'close', 'volume', 'volCcy', 'volCcyQuote', 'confirm']
        df = pd.DataFrame(data, columns=columns)
        
        # 转换数据类型
        numeric_columns = ['open', 'high', 'low', 'close', 'volume', 'volCcy', 'volCcyQuote']
        for col in numeric_columns:
            df[col] = pd.to_numeric(df[col])
        
        # 转换时间戳为datetime - 先转换为数值类型，避免FutureWarning
        df['timestamp'] = pd.to_numeric(df['timestamp'])
        df['timestamp'] = pd.to_datetime(df['timestamp'], unit='ms')
        
        # 转换为UTC+8时区
        df['timestamp'] = df['timestamp'].dt.tz_localize('UTC').dt.tz_convert(self.timezone)
        
        # 设置时间戳为索引并排序
        df.set_index('timestamp', inplace=True)
        df.sort_index(inplace=True)
        
        # 重命名列以符合mplfinance要求
        df.rename(columns={
            'open': 'Open',
            'high': 'High',
            'low': 'Low',
            'close': 'Close',
            'volume': 'Volume'
        }, inplace=True)
        
        return df
    
    async def fetch_historical_data(self):
        """
        获取历史K线数据
        
        返回:
            pandas.DataFrame: 历史K线数据
        """
        print(f"获取 {self.symbol} 的历史K线数据...")
        data = await self.market.get_candlesticks(self.symbol, self.interval, self.limit)
        if data:
            print(f"获取到 {len(data)} 条历史K线数据")
            self.df = self.process_kline_data(data)
            self.calculate_indicators()
            return self.df
        else:
            print("获取历史K线数据失败")
            return None
    
    def calculate_indicators(self):
        """计算技术指标"""
        if self.df is not None and len(self.df) > 0:
            # 计算移动平均线
            self.df['MA5'] = self.df['Close'].rolling(window=5).mean()
            self.df['MA10'] = self.df['Close'].rolling(window=10).mean()
            self.df['MA20'] = self.df['Close'].rolling(window=20).mean()
            
            # 计算布林带 (20,2)
            self.df['MA20'] = self.df['Close'].rolling(window=20).mean()
            self.df['UpperBand'] = self.df['MA20'] + 2 * self.df['Close'].rolling(window=20).std()
            self.df['LowerBand'] = self.df['MA20'] - 2 * self.df['Close'].rolling(window=20).std()
    
    async def start_websocket(self):
        """启动WebSocket连接，订阅K线数据"""
        # 根据OKX API文档，正确的频道名称格式
        channel_map = {
            "1m": "candle1m",
            "3m": "candle3m",
            "5m": "candle5m",
            "15m": "candle15m",
            "30m": "candle30m",
            "1H": "candle1H",
            "2H": "candle2H",
            "4H": "candle4H",
            "6H": "candle6H",
            "12H": "candle12H",
            "1D": "candle1D",
            "1W": "candle1W",
            "1M": "candle1M"
        }
        
        # 确保使用正确的频道名称
        channel = channel_map.get(self.interval, f"candle{self.interval}")
        
        # 构建订阅参数，确保格式正确
        candle_args = [{
            "channel": channel,  # 使用映射后的时间间隔
            "instId": self.symbol
        }]
        
        # 同时订阅tickers频道，以获取实时价格更新
        ticker_args = [{
            "channel": "tickers",
            "instId": self.symbol
        }]
        
        try:
            # 订阅K线频道
            await self.market.subscribe(candle_args, self.on_kline_message)
            print(f"已订阅K线: {self.symbol} {channel}")
            
            # 订阅tickers频道
            await self.market.subscribe(ticker_args, self.on_ticker_message)
            print(f"已订阅行情: {self.symbol}")
        except Exception as e:
            print(f"订阅数据失败: {e}")
            # 尝试重新连接
            print("尝试重新连接...")
            await asyncio.sleep(5)  # 等待5秒后重试
            await self.start_websocket()
    
    def on_kline_message(self, message):
        """
        处理K线WebSocket消息
        
        参数:
            message (str): WebSocket消息
        """
        try:
            data = json.loads(message)
            if 'data' in data:
                kline_data = data['data'][0]
                
                # 确保kline_data是列表类型
                if not isinstance(kline_data, list) or len(kline_data) < 6:
                    return
                
                # 解析K线数据
                timestamp_ms = pd.to_numeric(kline_data[0])
                timestamp = pd.to_datetime(timestamp_ms, unit='ms')
                timestamp = timestamp.tz_localize('UTC').tz_convert(self.timezone)
                
                open_price = float(kline_data[1])
                high_price = float(kline_data[2])
                low_price = float(kline_data[3])
                close_price = float(kline_data[4])
                volume = float(kline_data[5])
                volCcy = float(kline_data[6]) if len(kline_data) > 6 else 0
                
                # 创建新的K线数据行
                new_row = pd.DataFrame({
                    'Open': [open_price],
                    'High': [high_price],
                    'Low': [low_price],
                    'Close': [close_price],
                    'Volume': [volume],
                    'VolCcy': [volCcy]  # 添加成交额字段
                }, index=[timestamp])
                
                # 更新DataFrame
                self.update_dataframe(new_row)
                
                # 打印接收到的K线数据
                print(f"收到K线数据: {self.symbol} {timestamp} 开盘: {open_price} 最高: {high_price} 最低: {low_price} 收盘: {close_price} 成交量: {volume} 成交额: {volCcy}")
                
        except Exception as e:
            print(f"处理K线WebSocket消息时出错: {e}")
    
    def on_ticker_message(self, message):
        """
        处理Ticker WebSocket消息
        
        参数:
            message (str): WebSocket消息
        """
        try:
            data = json.loads(message)
            
            # 处理心跳消息或订阅确认消息
            if 'event' in data:
                return
                
            if 'data' in data:
                ticker_data = data['data'][0]
                
                # 确保ticker_data是字典类型
                if not isinstance(ticker_data, dict):
                    return
                
                # 获取当前时间作为时间戳
                now = datetime.now(self.timezone)
                
                # 根据K线周期取整时间
                if self.interval == '1m':
                    # 分钟级别K线，取整到分钟
                    timestamp = now.replace(second=0, microsecond=0)
                elif self.interval == '3m':
                    # 3分钟K线，取整到最近的3分钟
                    minute = (now.minute // 3) * 3
                    timestamp = now.replace(minute=minute, second=0, microsecond=0)
                elif self.interval == '5m':
                    # 5分钟K线，取整到最近的5分钟
                    minute = (now.minute // 5) * 5
                    timestamp = now.replace(minute=minute, second=0, microsecond=0)
                elif self.interval == '15m':
                    # 15分钟K线，取整到最近的15分钟
                    minute = (now.minute // 15) * 15
                    timestamp = now.replace(minute=minute, second=0, microsecond=0)
                elif self.interval == '30m':
                    # 30分钟K线，取整到最近的30分钟
                    minute = (now.minute // 30) * 30
                    timestamp = now.replace(minute=minute, second=0, microsecond=0)
                elif self.interval == '1H' or self.interval == '1h':
                    # 1小时K线，取整到小时
                    timestamp = now.replace(minute=0, second=0, microsecond=0)
                elif self.interval == '2H' or self.interval == '2h':
                    # 2小时K线，取整到最近的2小时
                    hour = (now.hour // 2) * 2
                    timestamp = now.replace(hour=hour, minute=0, second=0, microsecond=0)
                elif self.interval == '4H' or self.interval == '4h':
                    # 4小时K线，取整到最近的4小时
                    hour = (now.hour // 4) * 4
                    timestamp = now.replace(hour=hour, minute=0, second=0, microsecond=0)
                elif self.interval == '6H' or self.interval == '6h':
                    # 6小时K线，取整到最近的6小时
                    hour = (now.hour // 6) * 6
                    timestamp = now.replace(hour=hour, minute=0, second=0, microsecond=0)
                elif self.interval == '12H' or self.interval == '12h':
                    # 12小时K线，取整到最近的12小时
                    hour = (now.hour // 12) * 12
                    timestamp = now.replace(hour=hour, minute=0, second=0, microsecond=0)
                elif self.interval == '1D' or self.interval == '1d':
                    # 日K线，取整到天
                    timestamp = now.replace(hour=0, minute=0, second=0, microsecond=0)
                elif self.interval == '1W' or self.interval == '1w':
                    # 周K线，取整到周一
                    days_since_monday = now.weekday()
                    timestamp = (now - timedelta(days=days_since_monday)).replace(hour=0, minute=0, second=0, microsecond=0)
                elif self.interval == '1M' or self.interval == '1m':
                    # 月K线，取整到月初
                    timestamp = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
                else:
                    # 默认取整到分钟
                    timestamp = now.replace(second=0, microsecond=0)
                
                # 获取当前价格
                current_price = float(ticker_data.get('last', 0))
                
                # 如果当前时间戳在DataFrame中不存在，则创建新的K线
                if timestamp not in self.df.index:
                    # 创建新的K线数据行
                    new_row = pd.DataFrame({
                        'Open': [current_price],
                        'High': [current_price],
                        'Low': [current_price],
                        'Close': [current_price],
                        'Volume': [0],  # 初始成交量为0
                        'VolCcy': [0]   # 初始成交额为0
                    }, index=[timestamp])
                    
                    # 更新DataFrame
                    self.update_dataframe(new_row)
                    print(f"创建新的K线: {self.symbol} {timestamp} 开盘价: {current_price}")
                else:
                    # 更新现有K线的收盘价、最高价和最低价
                    self.df.at[timestamp, 'High'] = max(self.df.at[timestamp, 'High'], current_price)
                    self.df.at[timestamp, 'Low'] = min(self.df.at[timestamp, 'Low'], current_price)
                    self.df.at[timestamp, 'Close'] = current_price
                    
                    # 重新计算技术指标
                    self.calculate_indicators()
                
                # 更新图表 - 每次收到新的ticker数据都更新图表
                self.update_chart(0)
                
                print(f"收到Ticker数据: {self.symbol} {timestamp} 价格: {current_price}")
                
        except Exception as e:
            print(f"处理Ticker WebSocket消息时出错: {e}")
    
    def update_dataframe(self, new_data=None):
        """
        使用最新的K线数据更新DataFrame
        
        参数:
            new_data (pd.DataFrame): 最新的K线数据
        """
        if self.df is None:
            self.df = pd.DataFrame(columns=['Open', 'High', 'Low', 'Close', 'Volume', 'VolCcy'])
            return
            
        if new_data is None:
            return
            
        # 获取时间戳
        timestamp = new_data.index[0]
        
        # 检查是否已存在相同时间戳的数据
        if timestamp in self.df.index:
            # 更新现有数据，保留最高价和最低价
            self.df.loc[timestamp, 'Open'] = new_data.loc[timestamp, 'Open']
            self.df.loc[timestamp, 'High'] = max(self.df.loc[timestamp, 'High'], new_data.loc[timestamp, 'High'])
            self.df.loc[timestamp, 'Low'] = min(self.df.loc[timestamp, 'Low'], new_data.loc[timestamp, 'Low'])
            self.df.loc[timestamp, 'Close'] = new_data.loc[timestamp, 'Close']
            self.df.loc[timestamp, 'Volume'] = new_data.loc[timestamp, 'Volume']
            
            # 更新成交额，如果存在
            if 'VolCcy' in new_data.columns and 'VolCcy' in self.df.columns:
                self.df.loc[timestamp, 'VolCcy'] = new_data.loc[timestamp, 'VolCcy']
        else:
            # 添加新数据
            self.df = pd.concat([self.df, new_data])
            
        # 保持数据按时间排序
        self.df.sort_index(inplace=True)
        
        # 限制数据量
        if len(self.df) > self.limit:
            self.df = self.df.iloc[-self.limit:]
            
        # 重新计算技术指标
        self.calculate_indicators()
    
    def update_chart(self, frame):
        """
        更新图表
        
        参数:
            frame: 动画帧
        """
        try:
            if self.df is None or len(self.df) == 0:
                return
                
            # 清除当前图表
            plt.clf()
            
            # 添加技术指标
            apds = []
            
            # 添加移动平均线
            if 'MA5' in self.df.columns and not self.df['MA5'].isna().all():
                apds.append(mpf.make_addplot(self.df['MA5'], color=self.colors[0], width=0.7, label='MA5'))
            
            if 'MA10' in self.df.columns and not self.df['MA10'].isna().all():
                apds.append(mpf.make_addplot(self.df['MA10'], color=self.colors[1], width=0.7, label='MA10'))
            
            if 'MA20' in self.df.columns and not self.df['MA20'].isna().all():
                apds.append(mpf.make_addplot(self.df['MA20'], color=self.colors[2], width=0.7, label='MA20'))
            
            # 添加布林带
            if 'UpperBand' in self.df.columns and not self.df['UpperBand'].isna().all():
                apds.append(mpf.make_addplot(self.df['UpperBand'], color=self.colors[3], width=0.7, label='Upper BB'))
            
            if 'LowerBand' in self.df.columns and not self.df['LowerBand'].isna().all():
                apds.append(mpf.make_addplot(self.df['LowerBand'], color=self.colors[4], width=0.7, label='Lower BB'))
            
            # 绘制K线图
            self.candle_plots = mpf.plot(
                self.df,
                type='candle',
                style=self.style,
                title=f'{self.symbol} {self.interval} K线图',
                ylabel='价格',
                ylabel_lower='成交量',
                volume=True,
                addplot=apds,
                figscale=1.5,
                panel_ratios=(4, 1),
                datetime_format='%Y-%m-%d %H:%M',
                xrotation=15,
                returnfig=True,  # 返回图表对象
                block=False
            )
            
            # 添加图例
            plt.legend(loc='upper left')
            
            # 显示最新价格
            latest_price = self.df['Close'].iloc[-1]
            latest_time = self.df.index[-1].strftime('%Y-%m-%d %H:%M:%S')
            plt.title(f'{self.symbol} {self.interval} K线图 - 最新价格: {latest_price:.2f} - {latest_time}', color='black')
            
            # 添加鼠标悬浮提示
            self._add_hover_info()
            
            # 强制重绘图表
            plt.draw()
            plt.gcf().canvas.flush_events()
            
        except Exception as e:
            print(f"更新图表时出错: {e}")
    
    def _add_hover_info(self):
        """添加鼠标悬浮显示K线详细信息的功能"""
        # 如果已经存在cursor，先移除
        if self.cursor:
            try:
                self.cursor.remove()
            except:
                pass
        
        # 获取当前图表的所有轴对象
        if self.candle_plots:
            fig, axes = self.candle_plots
            
            # 创建一个新的cursor对象 - 确保使用所有可见的轴，并设置持久显示
            self.cursor = mplcursors.cursor(axes=axes, hover=True, highlight=True)
            
            # 设置鼠标悬浮事件的持久性
            self.cursor.connect("add", lambda sel: sel.annotation.set(visible=True))
            self.cursor.connect("remove", lambda sel: sel.annotation.set(visible=False))
            
            # 为鼠标悬浮事件添加回调函数
            @self.cursor.connect("add")
            def on_add(sel):
                # 获取选中点的索引
                try:
                    # 获取选中点的x坐标（日期索引）
                    x_date = sel.target[0]
                    
                    # 将x坐标转换为日期
                    date_num = mdates.num2date(x_date)
                    date_str = date_num.strftime('%Y-%m-%d %H:%M:%S')
                    
                    # 找到最接近的日期索引
                    closest_date = None
                    min_diff = float('inf')
                    
                    for date_idx in self.df.index:
                        # 转换为matplotlib日期数值
                        date_num_val = mdates.date2num(date_idx.to_pydatetime())
                        diff = abs(date_num_val - x_date)
                        
                        if diff < min_diff:
                            min_diff = diff
                            closest_date = date_idx
                    
                    if closest_date is not None:
                        # 获取K线数据
                        kline_data = self.df.loc[closest_date]
                        
                        # 计算涨跌幅和涨跌额
                        prev_idx = self.df.index.get_loc(closest_date) - 1
                        prev_close = self.df['Close'].iloc[prev_idx] if prev_idx >= 0 else kline_data['Open']
                        change = kline_data['Close'] - prev_close
                        change_percent = (change / prev_close * 100) if prev_close != 0 else 0
                        
                        # 计算振幅
                        amplitude = ((kline_data['High'] - kline_data['Low']) / prev_close * 100) if prev_close != 0 else 0
                        
                        # 格式化时间
                        time_str = closest_date.strftime('%Y-%m-%d %H:%M:%S')
                        
                        # 构建显示文本
                        text = (
                            f"时间: {time_str}\n"
                            f"开盘: {kline_data['Open']:.2f}\n"
                            f"最高: {kline_data['High']:.2f}\n"
                            f"最低: {kline_data['Low']:.2f}\n"
                            f"收盘: {kline_data['Close']:.2f}\n"
                            f"涨跌额: {change:.2f}\n"
                            f"涨跌幅: {change_percent:.2f}%\n"
                            f"振幅: {amplitude:.2f}%\n"
                            f"成交量: {kline_data['Volume']:.4f}"
                        )
                        
                        # 添加成交额信息（如果存在）
                        if 'VolCcy' in kline_data and not pd.isna(kline_data['VolCcy']) and kline_data['VolCcy'] > 0:
                            text += f"\n成交额: {kline_data['VolCcy']:.2f}"
                        
                        # 设置显示文本
                        sel.annotation.set_text(text)
                        
                        # 根据涨跌设置文本颜色
                        if kline_data['Close'] >= kline_data['Open']:
                            sel.annotation.get_bbox_patch().set(fc="red", alpha=0.8)
                        else:
                            sel.annotation.get_bbox_patch().set(fc="green", alpha=0.8)
                        
                        # 确保注释保持可见
                        sel.annotation.set_visible(True)
                    else:
                        sel.annotation.set_text("无法获取K线数据")
                        sel.annotation.get_bbox_patch().set(fc="gray", alpha=0.8)
                except Exception as e:
                    print(f"处理鼠标悬浮事件时出错: {e}")
                    sel.annotation.set_text("无法获取K线数据")
                    sel.annotation.get_bbox_patch().set(fc="gray", alpha=0.8)
        else:
            # 如果没有candle_plots，使用当前图表
            self.cursor = mplcursors.cursor(hover=True, highlight=True)
            
            # 设置鼠标悬浮事件的持久性
            self.cursor.connect("add", lambda sel: sel.annotation.set(visible=True))
            self.cursor.connect("remove", lambda sel: sel.annotation.set(visible=False))
            
            @self.cursor.connect("add")
            def on_add(sel):
                try:
                    # 尝试获取x坐标
                    if hasattr(sel, 'target') and len(sel.target) >= 1:
                        x_date = sel.target[0]
                        
                        # 将x坐标转换为日期
                        date_num = mdates.num2date(x_date)
                        
                        # 找到最接近的日期索引
                        closest_date = None
                        min_diff = float('inf')
                        
                        for date_idx in self.df.index:
                            date_num_val = mdates.date2num(date_idx.to_pydatetime())
                            diff = abs(date_num_val - x_date)
                            
                            if diff < min_diff:
                                min_diff = diff
                                closest_date = date_idx
                        
                        if closest_date is not None:
                            # 获取K线数据
                            kline_data = self.df.loc[closest_date]
                            
                            # 计算涨跌幅和涨跌额
                            prev_idx = self.df.index.get_loc(closest_date) - 1
                            prev_close = self.df['Close'].iloc[prev_idx] if prev_idx >= 0 else kline_data['Open']
                            change = kline_data['Close'] - prev_close
                            change_percent = (change / prev_close * 100) if prev_close != 0 else 0
                            
                            # 计算振幅
                            amplitude = ((kline_data['High'] - kline_data['Low']) / prev_close * 100) if prev_close != 0 else 0
                            
                            # 格式化时间
                            time_str = closest_date.strftime('%Y-%m-%d %H:%M:%S')
                            
                            # 构建显示文本
                            text = (
                                f"时间: {time_str}\n"
                                f"开盘: {kline_data['Open']:.2f}\n"
                                f"最高: {kline_data['High']:.2f}\n"
                                f"最低: {kline_data['Low']:.2f}\n"
                                f"收盘: {kline_data['Close']:.2f}\n"
                                f"涨跌额: {change:.2f}\n"
                                f"涨跌幅: {change_percent:.2f}%\n"
                                f"振幅: {amplitude:.2f}%\n"
                                f"成交量: {kline_data['Volume']:.4f}"
                            )
                            
                            # 添加成交额信息（如果存在）
                            if 'VolCcy' in kline_data and not pd.isna(kline_data['VolCcy']) and kline_data['VolCcy'] > 0:
                                text += f"\n成交额: {kline_data['VolCcy']:.2f}"
                            
                            # 设置显示文本
                            sel.annotation.set_text(text)
                            
                            # 根据涨跌设置文本颜色
                            if kline_data['Close'] >= kline_data['Open']:
                                sel.annotation.get_bbox_patch().set(fc="red", alpha=0.8)
                            else:
                                sel.annotation.get_bbox_patch().set(fc="green", alpha=0.8)
                            
                            # 确保注释保持可见
                            sel.annotation.set_visible(True)
                        else:
                            sel.annotation.set_text("无法获取K线数据")
                            sel.annotation.get_bbox_patch().set(fc="gray", alpha=0.8)
                except Exception as e:
                    print(f"处理鼠标悬浮事件时出错: {e}")
                    sel.annotation.set_text("无法获取K线数据")
                    sel.annotation.get_bbox_patch().set(fc="gray", alpha=0.8)
    
    async def run(self):
        """运行K线图显示"""
        # 获取历史数据
        await self.fetch_historical_data()
        
        if self.df is None or len(self.df) == 0:
            print("没有K线数据，无法显示图表")
            return
        
        print("启动K线图显示...")
        
        # 创建图表（在主线程中）
        plt.ion()  # 打开交互模式
        self.fig, self.ax1 = plt.subplots(figsize=(12, 8), facecolor='white')
        self.ax1.set_facecolor('white')
        
        # 添加图表关闭事件处理
        self.fig.canvas.mpl_connect('close_event', self._on_close)
        
        # 初始绘制
        self.update_chart(0)
        
        # 启动WebSocket连接
        ws_task = asyncio.create_task(self.start_websocket())
        
        # 创建动画
        self.ani = FuncAnimation(
            self.fig,
            self.update_chart,
            interval=1000,  # 每秒更新一次
            blit=False,  # 禁用blitting以确保完整重绘
            cache_frame_data=False  # 禁用缓存以确保每次都重新绘制
        )
        
        # 显示图表（阻塞模式）
        plt.show(block=True)
        
        # 如果plt.show()返回，说明图表被关闭
        print("图表已关闭，程序即将退出...")
        self.stop()
        
        # 确保程序退出
        import sys
        sys.exit(0)
    
    def _on_close(self, event):
        """图表关闭事件处理函数"""
        print("图表已关闭，程序即将退出...")
        self.stop()
        # 确保程序退出
        import sys
        sys.exit(0)
    
    def stop(self):
        """停止K线图显示"""
        # 停止动画
        if self.ani:
            self.ani.event_source.stop()
        
        # 关闭WebSocket连接
        if hasattr(self.market, 'ws') and self.market.ws:
            try:
                asyncio.create_task(self.market.ws.close())
                print("已关闭WebSocket连接")
            except:
                pass
        
        # 确保图表已关闭
        if self.fig is not None:
            plt.close(self.fig)
            self.fig = None
            self.ax1 = None

async def main():
    """主函数"""
    chart = OKXMplfinanceChart(symbol="BTC-USDT-SWAP", interval="1m", limit=100, proxy="http://localhost:10809")
    await chart.run()

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("程序已停止")
    except Exception as e:
        print(f"发生错误: {e}") 